/*-
 * Copyright (c) 2007-2008 SINA Corporation, All Rights Reserved.
 * Author: Zhang shuo <zhangshuo@staff.sina.com.cn>
 * A gossip implementation.
 *
 * A red-black tree is used to store the node of peers. the key is peer id.
 * when a registeration command is received, new registered peer will add 
 * into the tree.
 *
 * whatever command received, system trys to gossip them to some other peers,
 * if the peer do not response, we sub its remove_flag, when remove_flag is 0,
 * the peer will be put into a Singly-linked tail queue and then remove it from
 * the tree, free the space. After that you should scan the removing peer list,
 * and gossip the disreg message refer to what listed.
 *
 */

#include <stdio.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <sys/time.h>
#include <sys/ioctl.h>
#include <sys/select.h>
#include <signal.h>
#include <errno.h>
#include <arpa/inet.h>


#include "gossip.h"


static void gossip_req_resp(GOSSIP_HANDLER *handler, PEER_NODE *peer_node, 
            int node_count, char *data, size_t len);


/* initialize the tree root and then declare the functions of RB-tree */
RB_PROTOTYPE_STATIC(RB_TREE, _gossip_node, node_entries, id_cmp);

/* the function of RB-tree to compare to keys */
static int
id_cmp(struct _gossip_node *a, struct _gossip_node *b)
{
	int diff = a->r_peerid - b->r_peerid;

	if (diff == 0)
		return (0);
	else if (diff < 0)
		return (-1);
	else
		return (1);
}

/* genrate the functions needed byu RB-tree */
RB_GENERATE_STATIC(RB_TREE, _gossip_node, node_entries, id_cmp);


//STAILQ_INIT(&my_http_req->patch_head);

static LOCAL_PEER _local_peer;
static GOSSIP_HANDLER _handler;

/* 
 * initialize the gossip module
 * notice that gossip_init is not thread safe
 * there is no requirment to init the gossip in every thread,
 * just initialize at the very start.
 */
GOSSIP_HANDLER *
gossip_init(char *init_host, int init_port, char *current_host, 
						int current_port, int default_ttl, int peer_id)
{
	int res;
    LOCAL_PEER *local_peer;
    GOSSIP_HANDLER *handler;

    local_peer = &_local_peer;
    handler = &_handler;

    if (init_host == NULL || current_host == NULL)
        return NULL;
    
    /* initialize the local peer info */
    memset(local_peer->l_host, '\0', GSP_HOST_NAME_LEN + 1);
    memmove(local_peer->l_host, current_host, strlen(current_host));
    local_peer->l_port = current_port;
    local_peer->l_ttl = default_ttl;
    local_peer->l_peerid = peer_id;
    
    memset(local_peer->init_host, '\0', GSP_HOST_NAME_LEN + 1);
    memmove(local_peer->init_host, init_host, strlen(init_host));
    local_peer->init_port = init_port;

    /* initialize the handler */
    handler->localpeer = local_peer;
    handler->last_command = 0;

    /* init the staiq link list */
    STAILQ_INIT(&handler->stq_list);
    
    /* initialize read-write lock */
    res = pthread_rwlock_init(&handler->tree_lock, NULL);
    res = pthread_rwlock_init(&handler->stq_lock, NULL);

    if (res != 0)
        return NULL;

    return handler;
}


int 
gossip_addpeer(GOSSIP_HANDLER *handler, char *r_host, int r_port, int r_peerid)
{
    PEER_NODE *peer_node;
    PEER_NODE *t_peer_node;

    /* create a new node */
    peer_node = (PEER_NODE *)GSP_MALLOC(sizeof(PEER_NODE));

    if (peer_node == NULL || handler == NULL || r_host == NULL)
        return GSP_PEER_ADD_FAILED;

    /* initialize the node */
    memset(peer_node->r_host, '\0', GSP_HOST_NAME_LEN + 1);
    memmove(peer_node->r_host, r_host, strlen(r_host));
    peer_node->r_port = r_port;
    peer_node->r_peerid = r_peerid;
    peer_node->last_cmd_id = 0;
    peer_node->remove_flag = GSP_TEST_TIMES;

    /* insert the peer info to the red-black tree */
    pthread_rwlock_wrlock(&handler->tree_lock);;
    t_peer_node = RB_INSERT(RB_TREE, &handler->rb_tree, peer_node);
    pthread_rwlock_unlock(&handler->tree_lock);

    /* red-black tree returns NULL when success */
    if (t_peer_node != NULL)
    {
        return GSP_PEER_ADD_FAILED;
    }

#ifdef GSP_DEBUG
    printf("(msg) [%d,%s:%d] ADD PEER OK.\n",handler->localpeer->l_peerid, handler->localpeer->l_host, 
                                                handler->localpeer->l_port);
#endif

    return GSP_PEER_ADD_OK;
}


/*
 * this funciton removes the specific peer.
 */
int 
gossip_rmpeer(GOSSIP_HANDLER *handler, int r_peerid, int addto_list)
{
    PEER_NODE temp_node;
    PEER_NODE *removed_node;
    PEER_NODE *node_find;
 
    /* create a node for RB-tree to search */
    temp_node.r_peerid = r_peerid;

    /* remove the peer from the tree */
    pthread_rwlock_wrlock(&handler->tree_lock);
    node_find = RB_FIND(RB_TREE, &handler->rb_tree, &temp_node);

    if (node_find != NULL)
    {
        removed_node = RB_REMOVE(RB_TREE, &handler->rb_tree, node_find);
    }
    else
    {
        removed_node = NULL;
    }

    pthread_rwlock_unlock(&handler->tree_lock);

    /* free the space allocated for the peer */
    if (removed_node == NULL)
        return GSP_PEER_REMOVE_FAILED;
    else if (addto_list == 1)
    {
        /* move the removing peer to disreg gossip list */

        /* the function of move the peer to deleting list is forbidden */

        /*
        pthread_rwlock_wrlock(&handler->stq_lock);
        STAILQ_INSERT_TAIL(&handler->stq_list, removed_node, stq_entry); 
        pthread_rwlock_unlock(&handler->stq_lock);
        */
    }
    else
    {
        GSP_FREE(removed_node);
    }

    return GSP_PEER_REMOVE_OK;
}
 
/* this function is not in use now */
void 
gossip_disreg(GOSSIP_HANDLER *handler)
{
    uint64_t msg_id;
    PEER_NODE *peer_node;
    PEER_NODE *t_peer_node;
    char msg_buf[GSP_BUF_LEN + 1];

    FILE *file;
    uint64_t rand_num;

    file = fopen("/dev/random", "r");
    fread(&rand_num, sizeof(uint64_t), 1, file);
    fclose(file);

    msg_id = rand_num;
    
    pthread_rwlock_wrlock(&handler->stq_lock);
    /* scan the link list and gossip the message */
    STAILQ_FOREACH_SAFE(peer_node, &handler->stq_list, stq_entry, t_peer_node)
    {
        #ifdef DEBUG
        printf("(msg)[%d,%s:%d send PEER message:\n]", 
                            handler->localpeer->l_peerid, 
                            handler->localpeer->l_host, 
                            handler->localpeer->l_port)
        #endif

        /* generate the disreg message */
        snprintf(msg_buf, GSP_BUF_LEN, "gsp_disreg %d %d %lu %d\r\n",
                    peer_node->r_peerid, handler->localpeer->l_ttl, msg_id, handler->localpeer->l_peerid);
        
        gossip_route(handler, msg_buf, strlen(msg_buf), msg_id, handler->localpeer->l_peerid, 0);
        STAILQ_REMOVE(&handler->stq_list, peer_node, _gossip_node, stq_entry);
        GSP_FREE(peer_node);
    }
    pthread_rwlock_unlock(&handler->stq_lock);
}  
 

/* 
 * this funciton implement an route method. it route specific message to random peers.
 * functions accepts 'msglen' bytes 'msg'. it will not send the message back to the 
 * just-now sender peer and the peer which generate the message.
 * it will select at most GSP_HOST_LIMIT different peers, and then call gossip_req_resp funciton.
 */
int 
gossip_route(GOSSIP_HANDLER *handler, char *msg, size_t msglen, uint64_t msgid, 
            int src_peerid, int sender_id)
{
    int max_peerid;
    int routing_id;
    int peer_count = 0;
    int i, j = 0;
    PEER_NODE *peer_node;
    PEER_NODE *node_find;
    PEER_NODE selected_node[GSP_HOST_LIMIT];
    PEER_NODE temp_node;
    char is_break_loop = 0;
    
    if (pthread_rwlock_rdlock(&handler->tree_lock) != 0)
        return GSP_ROUTE_FAILED;

#ifdef GSP_DEBUG
    PEER_NODE *ttt;

    printf("---------------------------------------------\n");
    RB_FOREACH(ttt, RB_TREE, &handler->rb_tree) 
    {
        printf(
               "[%d,%s:%d - rmFLAG=%d, lastCMD=%lu]\n", 
               ttt->r_peerid, ttt->r_host, ttt->r_port, ttt->remove_flag, ttt->last_cmd_id);
    }
    printf("---------------------------------------------\n");
#endif
    
    /* select at most 'GSP_HOST_LIMIT' hosts to send */
    peer_node = RB_MAX(RB_TREE, &handler->rb_tree);

    if (peer_node == NULL)
    {
        pthread_rwlock_unlock(&handler->tree_lock);
        return GSP_ROUTE_FAILED;
    }

    max_peerid = peer_node->r_peerid;
     
    while(peer_count < GSP_HOST_LIMIT)
    {
      
        #ifdef GSP_DEBUG_ADVANCED
        printf("j = %d\n", j);
        #endif

        is_break_loop = 0;
      
        /* max retry times */
        if (j++ >= GSP_HOST_LIMIT * 4)
        {
            break;
        }
        
        /* compute the random id */
        routing_id = (rand() % max_peerid) + 1;

        #ifdef GSP_DEBUG_ADVANCED
        printf("max_peerid is %d. routing_id = %d\n", max_peerid + 1, routing_id);
        #endif

        /* do not send the msg back to the peer generate the msg */
        if (src_peerid == routing_id || sender_id == routing_id)
        {
            #ifdef GSP_DEBUG_ADVANCED
            printf("skip msg src peer or sender.\n");
            #endif
            is_break_loop = 1;
        }
        else
        {
            for (i = 0; i < peer_count; i++)
            {   
                /* filtrate the repeated 'routing_id' */  
                if (selected_node[i].r_peerid == routing_id)
                {
                    #ifdef GSP_DEBUG_ADVANCED
                    printf("peer repeated!\n");
                    #endif
                    is_break_loop = 1;
                }
                else
                {
                    /* if the last_cmd_id == msgid, skip */
                    if (selected_node[i].last_cmd_id == msgid)
                    {
                        #ifdef GSP_DEBUG_ADVANCED
                        printf("msg repeated! msgid=%d\n", msgid);
                        #endif
                        is_break_loop = 1;
                    }
                }
            }
        }

        if (is_break_loop) continue;

        /* find the peer via id */
        temp_node.r_peerid = routing_id;
        node_find = RB_FIND(RB_TREE, &handler->rb_tree, &temp_node);

        /* found (otherwise maybe id is not continuous) */
        if (node_find != NULL)
        {
            #ifdef GSP_DEBUG_ADVANCED
            printf("node_find(last_cmd_id):%d\n", node_find->last_cmd_id);
            printf("node_find(r_peerid):%d\n", node_find->r_peerid);
            #endif
          
            memmove(&selected_node[peer_count], node_find, sizeof(PEER_NODE));  
          
            /*
             * here try to modify the value under a read lock.
             * but it is nothing serious, it is only be used in filtrating the
             * repeated command. little mistake is allowed.
             */

            node_find->last_cmd_id = msgid;    /* save the message id */
            peer_count++;
        }
        else
        {
            #ifdef GSP_DEBUG_ADVANCED
            printf("peer id not found,\n");
            #endif
        }
    }
    
    /* unlock */
    pthread_rwlock_unlock(&handler->tree_lock);
    
    /* send the command one by one and then wait for response */
    gossip_req_resp(handler, selected_node, peer_count, msg, msglen);

    return GSP_ROUTE_OK;
}


/*
 * connect to specific peers, send message to connected peers and then waiting for the peer's resposne.
 * this function accept a peer node array. the length of the array is indicated by
 * 'node_count'. 'peer_node' is the pointer of the node of red-black tree.
 * finally the 'data' is the message to be send while 'len' is the length of the message.
 */
static void 
gossip_req_resp(GOSSIP_HANDLER *handler, PEER_NODE *peer_node, int node_count,
                char *data, size_t len)
{
    int i;
    int max_fd = 0;
    int socketfd[GSP_HOST_LIMIT][2];
    int waiting_peers = 0;
    int waiting_times = 0;
    int res;
    fd_set rsetall;
    fd_set rset;
    char readbuf[GSP_BUF_LEN];
    PEER_NODE temp_node;
    PEER_NODE *t_peer_node;

    struct timeval timeout;

    FD_ZERO(&rsetall);       		/* clear the read set of function select */
    for (i = 0; i < GSP_HOST_LIMIT; i++) socketfd[i][0] = -1;

    #ifdef GSP_DEBUG
    printf("(msg) finaly selected [%d] hosts.\n", node_count);
    #endif
    
    /* connect the selected peer */
    for (i = 0; i < node_count; i++)
    {
        socketfd[i][0] = gossip_connect(peer_node[i].r_host, peer_node[i].r_port);
        socketfd[i][1] = peer_node[i].r_peerid;

        if (socketfd[i][0] > max_fd) max_fd = socketfd[i][0];

        #ifdef GSP_DEBUG
        printf(">>> message send to: [%d,%s:%d]\n", peer_node[i].r_peerid, peer_node[i].r_host, peer_node[i].r_port);
        printf("\t%s\n", data);
        #endif
        
        if (socketfd[i][0] < 0)
        {
            #ifdef GSP_DEBUG
            printf("\t>>> connect %s[%d] failed!\n", peer_node[i].r_host, peer_node[i].r_port);
            #endif

            temp_node.r_peerid = socketfd[i][1];
           
            /* write lock */
            pthread_rwlock_wrlock(&handler->tree_lock);

            /* find and reset the peer's sub the remove_flag */
            t_peer_node = RB_FIND(RB_TREE, &handler->rb_tree, &temp_node);
            if (t_peer_node == NULL) 
            {
                pthread_rwlock_unlock(&handler->tree_lock);
                continue;
            }
            t_peer_node->remove_flag--;

            /* unlock */
            pthread_rwlock_unlock(&handler->tree_lock);

            /* if the remove_flag less than 0, we remove it */
            if (t_peer_node->remove_flag <= 0)
            {
                #ifdef GSP_DEBUG_ADVANCED
                printf(">>> peer removed:%s[%d]\n", peer_node[i].r_host, peer_node[i].r_port);  
                #endif
                gossip_rmpeer(handler, t_peer_node->r_peerid, 0);
            }
        }
        else
        {
            gossip_writen(socketfd[i][0], data, len);
            FD_SET(socketfd[i][0], &rsetall); 
            waiting_peers++;
        }
    }

    #ifdef GSP_DEBUG
    printf("(msg) waiting [%d] peers to response...\n", waiting_peers);
    #endif

    /*
     * this part waits for the peer response.
     * N peers were selected above, and waiting_times & waiting_peers are equal to N.
     * waiting_times decrease after each excuting of function SELECT.
     * waiting_peers decrease after receiving a successful response.
     * All this is to realize: either SELECT excuting N times or receiving N successful
     * responses, the loop will be exit.
     */

    waiting_times = waiting_peers;

    while (waiting_peers > 0 && waiting_times > 0)
    {
        rset = rsetall;  
        
        #ifdef GSP_DEBUG
        printf(">>> entering select loop...\n");
        #endif
        
        timeout.tv_sec = PEER_TIMEOUT_SEC;
        timeout.tv_usec = PEER_TIMEOUT_USEC;

        select(max_fd + 1, &rset, NULL, NULL, &timeout);

        waiting_times--;

        for (i = 0; i < node_count; i++)
        {
            /* get the peer id which responsed */
            temp_node.r_peerid = socketfd[i][1];
          
            if (socketfd[i][0] < 0)
            { 
                continue;
            }
            
            if (FD_ISSET(socketfd[i][0], &rset))
            {
                res = read(socketfd[i][0], readbuf, 6);
                if (strncmp("DONE\r\n", readbuf, 6) == 0)
                {
                    #ifdef GSP_DEBUG
                    printf("\tresponse from peer ID: [%d] OK.\n", socketfd[i][1]); 
                    #endif
                  
                   /* 
                     * if remote peer do not response DONE\r\n.
                     * we think it is in some error state. we do not
                     * reset the 'remove_flag'.
                     */

                    waiting_peers--;

                    /* write lock */
                    pthread_rwlock_wrlock(&handler->tree_lock);

                    /* find and reset the peer's remove_flag to GSP_TEST_TIMES */
                    t_peer_node = RB_FIND(RB_TREE, &handler->rb_tree, &temp_node);
                    if (t_peer_node != NULL)
                    {
                        t_peer_node->remove_flag = GSP_TEST_TIMES;
                    }

                    /* unlock */
                    pthread_rwlock_unlock(&handler->tree_lock);
                }
                else
                {
                    #ifdef GSP_DEBUG
                    printf("\tresponse from peer ID: [%d] ERROR. time out (PEER_TIMEOUT_SEC seconds).\n", socketfd[i][1]); 
                    #endif

                    do{}while(0);
                }

                FD_CLR(socketfd[i][0], &rsetall);
                close(socketfd[i][0]);
				socketfd[i][0] = -1; 
            }
        }
    }

    /* scan the list, and close the NON-RESPONSE peer */
    for (i = 0; i < node_count; i++)
    {
        if (socketfd[i][0] >= 0)
        { 
             close(socketfd[i][0]);
        }
    }

    return;
}


/*
 * clean the gossip resource.
 * it should be called at the very end of your programme
 */
int 
gossip_clean(GOSSIP_HANDLER *handler)
{
    PEER_NODE *var;
    PEER_NODE *nxt;
    
    if (handler == NULL)
        return -1;
    
    /* destory the lock */
    pthread_rwlock_destroy(&handler->tree_lock);

    /* free each node */
    for (var = RB_MIN(RB_TREE, &handler->rb_tree); var != NULL; var = nxt) 
    {
        nxt = RB_NEXT(RB_TREE, &handler->rb_tree, var);
        RB_REMOVE(RB_TREE, &handler->rb_tree, var);
        GSP_FREE(var);
    }
    
    return 1;
}


/*
 * connect to remote server
 */
int 
gossip_connect(char *host, int port)
{
    int sockfd;
    struct sockaddr_in	servaddr;
    
    sockfd = socket(AF_INET, SOCK_STREAM, 0);

    bzero(&servaddr, sizeof(servaddr));
    servaddr.sin_family      = AF_INET;
    servaddr.sin_addr.s_addr = inet_addr(host);
    servaddr.sin_port        = htons(port);
	    
    if (connect(sockfd, (struct sockaddr *)&servaddr, sizeof(servaddr)) < 0)
    {
		close(sockfd);
        return -1;
    }

    return sockfd;
}


/* 
 * write N bytes data to specific sockfd.
 * used in block-IO.
 */
ssize_t	
gossip_writen(int fd, const void *vptr, ssize_t n)
{
	ssize_t		nleft;
    ssize_t		nwritten;
	const char	*ptr;

	ptr = vptr;
	nleft = n;
	while (nleft > 0) 
    {
		if ((nwritten = write(fd, ptr, nleft)) <= 0) 
        {
			if (errno == EINTR)
				nwritten = 0;		/* and call write() again */
			else
				return(-1);			/* error */
		}

		nleft -= nwritten;
		ptr   += nwritten;
	}
	return n;
}



